from pysnmp.hlapi import *
from datetime import datetime

class SNMPCollector:
    def __init__(self, host):
        self.host = host
        self.community = host.snmp_community
        self.port = host.snmp_port
        self.version = host.snmp_version

    def get_cpu_usage(self):
        try:
            # CPU负载 OID (使用标准MIB-2)
            cpu_load_oid = '1.3.6.1.4.1.2021.11.9.0'  # ssCpuIdle
            result = self._get_snmp_value(cpu_load_oid)
            if result is not None:
                return 100 - float(result)  # 转换空闲率为使用率
        except Exception as e:
            print(f"获取CPU使用率失败: {str(e)}")
        return None

    def get_memory_usage(self):
        try:
            # 内存相关OID
            total_memory_oid = '1.3.6.1.4.1.2021.4.5.0'  # memTotalReal
            free_memory_oid = '1.3.6.1.4.1.2021.4.6.0'   # memAvailReal
            
            total = float(self._get_snmp_value(total_memory_oid))
            free = float(self._get_snmp_value(free_memory_oid))
            
            if total and free:
                return ((total - free) / total) * 100
        except Exception as e:
            print(f"获取内存使用率失败: {str(e)}")
        return None

    def get_disk_usage(self):
        try:
            # 磁盘使用率OID (根文件系统)
            disk_used_oid = '1.3.6.1.4.1.2021.9.1.9.1'   # dskPercent
            result = self._get_snmp_value(disk_used_oid)
            if result is not None:
                return float(result)
        except Exception as e:
            print(f"获取磁盘使用率失败: {str(e)}")
        return None

    def _get_snmp_value(self, oid):
        try:
            error_indication, error_status, error_index, var_binds = next(
                getCmd(SnmpEngine(),
                      CommunityData(self.community, mpModel=0 if self.version == '1' else 1),
                      UdpTransportTarget((self.host.ip_address, self.port)),
                      ContextData(),
                      ObjectType(ObjectIdentity(oid)))
            )

            if error_indication:
                print(f"SNMP错误: {error_indication}")
                return None
            elif error_status:
                print(f"SNMP错误状态: {error_status}")
                return None
            else:
                return var_binds[0][1]
        except Exception as e:
            print(f"SNMP查询失败: {str(e)}")
            return None

    def collect_all(self):
        """收集所有SNMP数据并更新主机信息"""
        try:
            self.host.snmp_cpu_usage = self.get_cpu_usage()
            self.host.snmp_memory_usage = self.get_memory_usage()
            self.host.snmp_disk_usage = self.get_disk_usage()
            self.host.snmp_last_update = datetime.now()
            self.host.save()
        except Exception as e:
            print(f"SNMP数据收集失败: {str(e)}")
            raise 